import pdfplumber
import re
import pandas as pd
from pathlib import Path


def extract_income_statement(pdf_path):
    with pdfplumber.open(pdf_path) as pdf:
        # 改进的章节定位方式
        financial_section = []
        in_financial_section = False
        table_found = False

        for page in pdf.pages:
            text = page.extract_text()
            if not text:
                continue

            # 定位财务报告章节
            if "第十节" in text and "财务报告" in text:
                in_financial_section = True
                print(f"找到财务报告章节在页面 {page.page_number}")

            if in_financial_section:
                # 寻找合并利润表
                if "合并利润表" in text:
                    print(f"发现合并利润表在页面 {page.page_number}")
                    table_found = True

                if table_found:
                    # 改进的表格提取逻辑
                    tables = page.extract_tables({
                        "vertical_strategy": "text",
                        "horizontal_strategy": "text"
                    })

                    for table in tables:
                        cleaned_table = []
                        for row in table:
                            clean_row = [cell.replace('\n', ' ') if cell else '' for cell in row]
                            if any(cell.strip() for cell in clean_row):
                                cleaned_table.append(clean_row)

                        # 通过特征识别利润表
                        if len(cleaned_table) >= 5 and "营业总收入" in cleaned_table[0][0]:
                            print("找到有效的利润表结构")
                            df = pd.DataFrame(cleaned_table[1:], columns=cleaned_table[0])
                            return df

            # 结束条件
            if table_found and "法定代表人" in text:
                break

        raise ValueError("未能提取合并利润表")


def main():
    pdf_dir = Path(r"D:\PycharmProjects\PythonProject\阳光电源年报")

    for pdf_file in pdf_dir.glob("*.pdf"):
        print(f"\n{'=' * 30}\n正在处理：{pdf_file.name}")

        try:
            df = extract_income_statement(pdf_file)

            # 数据清洗
            df = df.applymap(lambda x: x.strip() if isinstance(x, str) else x)
            df.replace(['', '—', '--'], pd.NA, inplace=True)
            df.dropna(how='all', inplace=True)

            # 保存结果
            output_path = pdf_dir / f"{pdf_file.stem}_合并利润表.xlsx"
            df.to_excel(output_path, index=False)
            print(f"成功保存：{output_path}")

        except Exception as e:
            print(f"处理失败：{str(e)}")
            continue


if __name__ == "__main__":
    main()